#include <WiFi.h>
#include <PubSubClient.h>
#include <ArduinoJson.h>
#include <Adafruit_NeoPixel.h>

const char* ssid      = "Shahmeer";
const char* password  = "redp5314";
const char* broker    = "192.168.128.176";

#define LED_PIN    17
#define LED_COUNT  3
#define THRESHOLD  15.0

Adafruit_NeoPixel strip(LED_COUNT, LED_PIN, NEO_GRB + NEO_KHZ800);
WiFiClient wifi;
PubSubClient mqtt(wifi);
bool shakeDetected = false;

void setStrip(uint8_t r, uint8_t g, uint8_t b) {
  for (int i = 0; i < strip.numPixels(); i++) {
    strip.setPixelColor(i, strip.Color(r, g, b));
  }
  strip.show();
}

void onMessage(char* topic, byte* payload, unsigned int length) {
  StaticJsonDocument<200> doc;
  deserializeJson(doc, payload, length);
  float ax = doc["ax"];
  float ay = doc["ay"];
  float az = doc["az"];
  float magnitude = sqrt(ax*ax + ay*ay + az*az);
  Serial.print("magnitude: ");
  Serial.println(magnitude);
  shakeDetected = (magnitude > THRESHOLD);
}

void connectWifi() {
  WiFi.begin(ssid, password);
  while (WiFi.status() != WL_CONNECTED) delay(500);
  Serial.println("WiFi connected");
  Serial.println(WiFi.localIP());
}

void connectMQTT() {
  setStrip(0, 0, 255);
  while (!mqtt.connected()) {
    if (mqtt.connect("esp32-subscriber")) {
      mqtt.subscribe("fab/imu");
      Serial.println("MQTT connected");
    } else {
      Serial.print("failed rc=");
      Serial.println(mqtt.state());
      delay(500);
    }
  }
}

void setup() {
  Serial.begin(115200);
  strip.begin();
  strip.show();
  setStrip(0, 0, 255);
  connectWifi();
  mqtt.setKeepAlive(60);
  mqtt.setServer(broker, 1883);
  mqtt.setCallback(onMessage);
  connectMQTT();
  setStrip(0, 255, 0);
}

void loop() {
  if (!mqtt.connected()) connectMQTT();
  mqtt.loop();

  static unsigned long lastUpdate = 0;
  if (millis() - lastUpdate > 50) {
    lastUpdate = millis();
    if (shakeDetected) {
      setStrip(255, 0, 0);
    } else {
      setStrip(0, 255, 0);
    }
  }
}